
from yacf.Frontend.C99 import c99_ast

from yacf.Backends.C99.Writers.C99Writer import CWriter

class OmpWriter(CWriter):
    """ Visitor which translates the IR to C/OpenMP.

    """
    def visit_OmpParallel(self, node, offset):
        self.write_blank();
        self.write(offset, node.name)
        self.write_blank();
        # self.write(offset, 'parallel')
        # self.write_blank();
        if node.clauses:
            for elem in node.clauses:
                self.visit(elem)
                self.write_blank();
        if node.stmt:
            self.writeLn(offset, "")
            self.visit(node.stmt)
        self.write_blank();

    def visit_OmpTargetDevice(self, node, offset):
        self.write_blank()
        self.write(offset, 'target')
        self.write_blank()
        self.write(offset, node.name)
        self.write(offset, '(')
        self.write(offset, node.device)
        self.write(offset, ')')

    def visit_OmpThreadPrivate(self, node, offset):
        self.write_blank();
        self.write(offset, node.name)
        self.write_blank()
        if node.identifiers:
            for elem in node.identifiers.params:
                self.write(offset, '(')
                self.visit(elem)
                self.write(offset, ')')
                self.write_blank()

    def visit_OmpFor(self, node, offset):
        self.write_blank();
        self.write(offset, node.name)
        self.write_blank();
        if node.clauses:
            for elem in node.clauses:
                self.visit(elem)
                self.write_blank();

        self.write_blank();
        # OmpFor always has an stmt 
        if node.stmt:
            self.writeLn(offset, " ")
            self.visit(node.stmt)
        self.write_blank();  


    def visit_OmpParallelFor(self, node, offset):
        self.write_blank();
        self.write(offset, node.name)
        self.write_blank();
        #=======================================================================
        # import pdb
        # pdb.set_trace()
        #=======================================================================
        if node.clauses:
            for elem in node.clauses:
                self.visit(elem)
                self.write_blank();

        self.write_blank();
        # OmpParallelFor always has an stmt 
        if node.stmt:
            self.writeLn(offset, " ")
            self.visit(node.stmt)
        self.write_blank();

    def visit_OmpSyncronization(self, node, offset):
        self.write_blank();
        self.write(offset, node.name)


    def visit_OmpClause(self, node, offset):
        # Handwrite the case of device clause
        if node.name == 'cuda':
            self.write(offset, 'device(' + node.name.lower() + ")")
        else:
            self.write(offset, node.name.lower())
        # TODO: Fix reduction parenthesis error
        if node.name == 'DEFAULT':
            self.write(0, '(');
            self.write(0, node.type)
            self.write(0, ')')

        if node.name == 'REDUCTION':
            self.write(0, '(');
            self.write(0, node.type)
            self.write(0, ':')

        if node.identifiers:
	        self.visit(node.identifiers)

        if node.name == 'REDUCTION':
            self.write(0, ')');
    
    ############################ llc nodes

    def visit_llcRegion(self, node, offset):
        self.write_blank();
        self.write(offset, node.name)
        self.write_blank();
        # self.write(offset, 'parallel')
        # self.write_blank();
        if node.clauses:
            for elem in node.clauses:
                self.visit(elem)
                self.write_blank();
        if node.stmt:
            self.writeLn(offset, "")
            self.visit(node.stmt)
        self.write_blank();


    def visit_llcFor(self, node, offset):
        self.write_blank();
        self.write(offset, node.name)
        self.write_blank();
        if node.clauses:
            for elem in node.clauses:
                self.visit(elem)
                self.write_blank();

        self.write_blank();
        # llcFor always has an stmt 
        if node.stmt:
            self.writeLn(offset, " ")
            self.visit(node.stmt)
        self.write_blank();  


    def visit_llcName(self, node, offset):
        self.write_blank();
        self.write(offset, "name( \"")
        self.write(offset, node.label)
        self.write(offset, "\")")
        


    def visit_llcResultItem(self, node, offset):
        self.write(0, '&')
        self.write(0, node.declaration.type.declname)
        self.visit(node.declaration)
        self.write(0, ',')
        self.visit(node.size)
        


    def visit_llcResult(self, node, offset):
        self.write_blank()
        self.write(offset, '#pragma llc result')
        self.write(0, '(')
        for elem in node.result_list:
            self.visit_llcResultItem(elem, 0)
            
        self.write(0, ')')
        self.writeLn(0, '\n')
        self.visit(node.stmt)

    def visit_llcNcResult(self, node, offset):
        self.write_blank()
        self.write(offset, '#pragma llc ncresult')
        self.write(0, '(')
        for elem in node.result_list:
            self.visit_llcResultItem(elem, 0)
            
        self.write(0, ')')
        self.writeLn(0, '\n')
        self.visit(node.stmt)


    def visit_CUDAKernelCall(self, node, offset = 0):
        self.visit_ID(node.name)
        self.write(offset, "<<<")
        self.visit(node.grid)
        self.write(offset, ",")
        self.visit(node.block)
        self.write(offset, ">>>")
        self.write(offset, "(")
        if node.args:
            self.visit_ExprList(node.args)
        else:
            self.write_blank()
        self.write(offset, ")")
        self.write_blank()

